/**
 * Copyright 厦门中软海晟信息技术有限公司 版权所有 违者必究 2019
 */
package com.opencvjava.dtk;

import com.google.common.collect.Lists;
import org.opencv.core.Core;
import org.opencv.core.Mat;
import org.opencv.core.Point;
import org.opencv.core.Point3;
import org.opencv.core.Rect;
import org.opencv.core.Scalar;
import org.opencv.core.Size;
import org.opencv.imgproc.Imgproc;
import org.springframework.stereotype.Service;

import java.util.ArrayList;
import java.util.HashMap;
import java.util.Iterator;
import java.util.List;
import java.util.Map;

import static com.opencvjava.support.util.CvUtils.imshow;
import static com.opencvjava.support.util.CvUtils.mat;
import static org.opencv.imgproc.Imgproc.*;
import static org.opencv.imgproc.Imgproc.MORPH_CLOSE;
import static org.opencv.imgproc.Imgproc.morphologyEx;

/**
 * @author : sunzb(sunzb@hsit.com.cn)
 * @date: 2019/1/22
 */
@Service
public class DtkTest {
    public void test() {
        // step1：加载图片
        Mat src = mat("dtk", "SheetStand_zkz.jpg");
        Mat template = mat("dtk", "Circle.jpg");
        if (src.empty() || template.empty()) {
            throw new RuntimeException("图片读取失败");
        }
        imshow("源图", src);
        imshow("模板", template);
        Mat copy = src.clone();
        Mat gray = new Mat();
        cvtColor(src, gray, COLOR_BGR2GRAY);
        imshow("灰度", gray);
        // step2：找出锚点
        List<Point> anchors = achors(copy, template);
        for (Point anchor : anchors) {
            rectangle(copy, new Rect((int) anchor.x, (int) anchor.y, template.cols(), template.rows()),
                    new Scalar(0, 0, 255), 2, 8, 0);
        }
        imshow("锚点", copy);
        // step3：根据锚点做仿射矫正
        // step4：找出各个区域
        Rois rois = rois(anchors, gray, template.cols(), template.rows());
        imshow("顶部", rois.getTop());
        imshow("底部", rois.getBottom());
        imshow("左边", rois.getLeft());
        imshow("右边", rois.getRight());
        imshow("内容", rois.getContentMat());
        // step5: 答题区域映射
        Mat binary = new Mat();
        threshold(rois.getTop(), binary, 100, 255, THRESH_BINARY);
        binary.copyTo(rois.getTop());
        imshow("顶部二值化", rois.getTop());
        threshold(rois.getBottom(), binary, 100, 255, THRESH_BINARY);
        binary.copyTo(rois.getBottom());
        imshow("底部二值化", rois.getBottom());
        threshold(rois.getLeft(), binary, 100, 255, THRESH_BINARY);
        imshow("左部二值化", binary);
        Mat element = getStructuringElement(CV_SHAPE_RECT,
                new Size(1, 10),
                new Point(0, 0));
        morphologyEx(binary, binary, MORPH_CLOSE, element);
        binary.copyTo(rois.getLeft());
        imshow("左部闭操作", rois.getLeft());
        threshold(rois.getRight(), binary, 100, 255, THRESH_BINARY);
        binary.copyTo(rois.getRight());
        imshow("右部二值化", rois.getRight());
        // step6: 元信息分块采集
        List<List<Point>> topMata = topAndBottomMata(src, rois.getTop(), rois.getTopP());
        List<List<Point>> bottomMata = topAndBottomMata(src, rois.getBottom(), rois.getBottomP());
        List<List<Point>> leftMata = leftOrRightMata(src, rois.getLeft(), rois.getLeftP());
        List<List<Point>> rightMata = leftOrRightMata(src, rois.getRight(), rois.getRightP());
        imshow("答案定位元信息", src);
        // 准考证号
        Map<Integer, List<Rect>> zkz = zkz(src, topMata, rightMata);
        StringBuilder zkzAnswer = new StringBuilder();
        Iterator<Map.Entry<Integer, List<Rect>>> iteratorZkz = zkz.entrySet().iterator();
        while (iteratorZkz.hasNext()) {
            Map.Entry<Integer, List<Rect>> next = iteratorZkz.next();
            List<Rect> rects = next.getValue();
            for (int i = 0; i < rects.size(); i++) {
                if (isAnswer(gray, rects.get(i), src)) {
                    zkzAnswer.append(i);
                    break;
                }
            }
        }
        imshow("准考证答案定位", src);
        System.out.println("准考证号是:" + zkzAnswer.toString());
    }

    private boolean isAnswer(Mat gray, Rect rect, Mat src) {
        rectangle(src, rect, new Scalar(255, 0, 0),1, 8, 0);
        Mat roi = gray.submat(rect);
        int total = (int) roi.total();
        byte[] data = new byte[total];
        roi.get(0, 0, data);
        int index, color, count = 0;
        for (int row = 0; row < roi.rows(); row++) {
            for (int col = 0; col < roi.cols(); col++) {
                index = row * roi.cols() + col;
                color = data[index] & 0xff;
                if (color < 10) {
                    count++;
                }
            }
        }
        System.out.println("count:" + count + ", total:" + total);
        return count != 0 && count > total / 2;
    }

    private Map<Integer, List<Rect>> zkz(Mat src, List<List<Point>> topMata, List<List<Point>> rightMata) {
        Point topLeft;
        Point topRight;
        Point rightTop;
        Point rightBottom;
        Map<Integer, List<Rect>> zkz = new HashMap<>();
        for (int i = 1; i < 10; i++) {
            topLeft = topMata.get(i).get(0);
            topRight = topMata.get(i).get(1);
            List<Rect> rects = Lists.newArrayList();
            zkz.put(i, rects);
            for (int j = 0; j < 10; j++) {
                rightTop = rightMata.get(j).get(0);
                rightBottom = rightMata.get(j).get(1);
                Rect rect = new Rect(new Point(topLeft.x, rightTop.y),
                        new Size(topRight.x - topLeft.x + 1, rightBottom.y - rightTop.y + 1));
                rects.add(rect);
            }
        }
        return zkz;
    }

    private List<List<Point>> leftOrRightMata(Mat src, Mat leftOrRight, Point leftOrRightP) {
        List<List<Point>> result = new ArrayList<>();
        int cols = leftOrRight.cols();
        int rows = leftOrRight.rows();
        byte[] data = new byte[(int) leftOrRight.total()];
        leftOrRight.get(0, 0, data);
        int index, foundCol = 0;
        for (int col = 0; col < cols; col++) {
            for (int row = 0; row < rows; row++) {
                index = row * cols + col;
                if (data[index] == 0) {
                    foundCol = col + 3;
                }
            }
            if (foundCol != 0) {
                break;
            }
        }
        boolean find = false;
        Point point1 = null;
        Point point2 = null;
        for (int row = 0; row < rows; row++) {
            index = row * cols + foundCol;
            if (!find && data[index] == 0) {
                find = true;
                point1 = new Point(leftOrRightP.x, leftOrRightP.y + row);
                line(src, point1,
                        new Point(leftOrRightP.x + 20, leftOrRightP.y + row),
                        new Scalar(0, 255, 0), 1, 8, 0);
            }
            if (find && data[index] == -1) {
                point2 = new Point(leftOrRightP.x, leftOrRightP.y + row - 1);
                line(src, point2,
                        new Point(leftOrRightP.x + 20, leftOrRightP.y + row - 1),
                        new Scalar(0, 255, 0), 1, 8, 0);
                find = false;
                result.add(Lists.newArrayList(point1, point2));
            }
        }
        return result;
    }

    private List<List<Point>> topAndBottomMata(Mat src, Mat topOrBottom, Point topOrBottomP) {
        List<List<Point>> result = new ArrayList<>();
        int cols = topOrBottom.cols();
        byte[] data = new byte[(int) topOrBottom.total()];
        topOrBottom.get(0, 0, data);
        int index, foundRow = 0;
        for (int row = 0; row < topOrBottom.rows(); row++) {
            for (int col = 0; col < cols; col++) {
                index = row * cols + col;
                if (data[index] == 0) {
                    foundRow = row + 3;
                    break;
                }
            }
            if (foundRow != 0) {
                break;
            }
        }
        boolean find = false;
        Point point1 = null;
        Point point2 = null;
        for (int col = 0; col < cols; col++) {
            index = foundRow * cols + col;
            if (!find && data[index] == 0) {
                find = true;
                point1 = new Point(topOrBottomP.x + col, topOrBottomP.y);
                line(src, point1,
                        new Point(topOrBottomP.x + col, topOrBottomP.y + 20),
                        new Scalar(0, 255, 0), 1, 8, 0);
            }
            if (find && data[index] == -1) {
                point2 = new Point(topOrBottomP.x + col - 1, topOrBottomP.y);
                line(src, point2,
                        new Point(topOrBottomP.x + col - 1, topOrBottomP.y + 20),
                        new Scalar(0, 255, 0), 1, 8, 0);
                find = false;
                result.add(Lists.newArrayList(point1, point2));
            }
        }
        return result;
    }

    private Rois rois(List<Point> anchors, Mat gray, int templateCols, int templateRows) {
        int tlx = (int) anchors.get(0).x + templateCols;
        int tly = (int) anchors.get(0).y + templateRows;
        int trx = (int) anchors.get(1).x;
        int width = trx - tlx;
        int height = (int) (anchors.get(2).y - (anchors.get(0).y + templateRows));
        Point topP = new Point(tlx, (int) anchors.get(0).y);
        Mat top = gray.submat(new Rect(topP, new Size(width, templateRows)));
        Point bottomP = new Point(tlx, (int) anchors.get(2).y);
        Mat bottom = gray.submat(new Rect(bottomP, new Size(width, templateRows)));
        Point rightP = new Point(trx, (int) anchors.get(0).y + templateRows);
        Mat right = gray.submat(new Rect(rightP, new Size(templateCols, height)));
        Point leftP = new Point((int) anchors.get(0).x, (int) anchors.get(0).y + templateRows);
        Mat left = gray.submat(new Rect(leftP, new Size(templateCols, height)));
        Point contentP = new Point(tlx, tly);
        Mat contentMat = gray.submat(new Rect(contentP, new Size(width, height)));
        return new Rois().setTop(top).setTopP(topP)
                .setBottom(bottom).setBottomP(bottomP)
                .setLeft(left).setLeftP(leftP)
                .setRight(right).setRightP(rightP)
                .setContentMat(contentMat);
    }

    private List<Point> achors(Mat src, Mat template) {
        int cx = src.cols() / 2;
        int cy = src.rows() / 2;
        List<Mat> rois = new ArrayList<>();
        rois.add(src.submat(new Rect(0, 0, cx, cy)));
        rois.add(src.submat(new Rect(cx, 0, cx, cy)));
        rois.add(src.submat(new Rect(0, cy, cx, cy)));
        rois.add(src.submat(new Rect(cx, cy, cx, cy)));
        List<Point> anchors = new ArrayList<>();
        Mat result = new Mat();
        int method = TM_CCOEFF_NORMED;
        for (int i = 0; i < rois.size(); i++) {
            Mat roi = rois.get(i);
            int offsetX = i % 2 == 0 ? 0 : cx;
            int offsetY = i < 2 ? 0 : cy;
            Imgproc.matchTemplate(roi, template, result, method);
            Core.MinMaxLocResult minMaxLocResult = Core.minMaxLoc(result);
            Point maxLoc = minMaxLocResult.maxLoc;
            Point minLoc = minMaxLocResult.minLoc;
            Point matchLoc;
            if (method == TM_SQDIFF || method == TM_SQDIFF_NORMED) {
                matchLoc = minLoc;
            } else {
                matchLoc = maxLoc;
            }
            anchors.add(new Point(offsetX + matchLoc.x, offsetY + matchLoc.y));
        }
        return anchors;
    }

    private static class Rois {
        Mat top;
        Point topP;
        Mat bottom;
        Point bottomP;
        Mat left;
        Point leftP;
        Mat right;
        Point rightP;
        Mat contentMat;
        Point contentP;

        public Point getTopP() {
            return topP;
        }

        public Rois setTopP(Point topP) {
            this.topP = topP;
            return this;
        }

        public Point getBottomP() {
            return bottomP;
        }

        public Rois setBottomP(Point bottomP) {
            this.bottomP = bottomP;
            return this;
        }

        public Point getLeftP() {
            return leftP;
        }

        public Rois setLeftP(Point leftP) {
            this.leftP = leftP;
            return this;
        }

        public Point getRightP() {
            return rightP;
        }

        public Rois setRightP(Point rightP) {
            this.rightP = rightP;
            return this;
        }

        public Point getContentP() {
            return contentP;
        }

        public Rois setContentP(Point contentP) {
            this.contentP = contentP;
            return this;
        }

        public Mat getTop() {
            return top;
        }

        public Rois setTop(Mat top) {
            this.top = top;
            return this;
        }

        public Mat getBottom() {
            return bottom;
        }

        public Rois setBottom(Mat bottom) {
            this.bottom = bottom;
            return this;
        }

        public Mat getLeft() {
            return left;
        }

        public Rois setLeft(Mat left) {
            this.left = left;
            return this;
        }

        public Mat getRight() {
            return right;
        }

        public Rois setRight(Mat right) {
            this.right = right;
            return this;
        }

        public Mat getContentMat() {
            return contentMat;
        }

        public Rois setContentMat(Mat contentMat) {
            this.contentMat = contentMat;
            return this;
        }
    }
}
